package zhihuilan.common.aspect;

import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpRequest;
import org.springframework.http.client.ClientHttpRequestExecution;
import org.springframework.http.client.ClientHttpRequestInterceptor;
import org.springframework.http.client.ClientHttpResponse;
import org.springframework.http.client.reactive.ClientHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.client.ClientRequest;
import org.springframework.web.reactive.function.client.ClientResponse;
import org.springframework.web.reactive.function.client.ExchangeFilterFunction;
import org.springframework.web.reactive.function.client.ExchangeFunction;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import zhihuilan.common.enums.Errors;
import zhihuilan.common.exception.BizException;

import java.io.IOException;
import java.nio.charset.StandardCharsets;

@Slf4j
@Component
public class LogClientHttpInterceptor implements ClientHttpRequestInterceptor, ExchangeFilterFunction {
    @Override
    public ClientHttpResponse intercept(HttpRequest request, byte[] body, ClientHttpRequestExecution execution) throws IOException {
        long startTime = System.currentTimeMillis();
        try {
            ClientHttpResponse result = execution.execute(request, body);
            if (result.getStatusCode().is2xxSuccessful()) {
                log.info("状态：调用外部成功，请求地址：{}:{}，请求头：{}，参数：{}，结果：{}，耗时：{}", request.getMethod(), request.getURI(), request.getHeaders(), new String(body, StandardCharsets.UTF_8), "", System.currentTimeMillis() - startTime);
            } else {
                log.error("状态：调用外部错误，请求地址：{}:{}，请求头：{}，参数：{}，结果：{}:{}，耗时：{}", request.getMethod(), request.getURI(), request.getHeaders(), new String(body, StandardCharsets.UTF_8), result.getStatusText(), "", System.currentTimeMillis() - startTime);
            }
            return result;
        } catch (Throwable e) {
            log.error("状态：调用外部异常，请求地址：{}:{}，请求头：{}，参数：{}，耗时：{}", request.getMethod(), request.getURI(), request.getHeaders(), new String(body, StandardCharsets.UTF_8), System.currentTimeMillis() - startTime, e);
            throw e;
        }
    }

    @Override
    public Mono<ClientResponse> filter(ClientRequest request, ExchangeFunction next) {
        long startTime = System.currentTimeMillis();
        StringBuilder requestBody = new StringBuilder();
        Mono<ClientResponse> response = next.exchange(buildTraceableRequest(request, requestBody))
                .flatMap(clientResponse -> {
                    if (!clientResponse.statusCode().is2xxSuccessful()) {
                        clientResponse.bodyToMono(DataBuffer.class).subscribe(dataBuffer -> {
                            log.error("状态：调用大模型异常，请求地址：{}:{}，请求头：{}，参数：{}，响应：{}:{}，耗时：{}", request.method(), request.url(), request.headers(), requestBody, clientResponse.statusCode(), dataBuffer.toString(StandardCharsets.UTF_8), System.currentTimeMillis() - startTime);
                        });
                        return Mono.error(new BizException(Errors.GPT_ERROR));
                    }
                    return Mono.just(clientResponse);
                })
                .doOnSuccess(result -> log.info("状态：调用大模型成功，请求地址：{}:{}，请求头：{}，参数：{}， 耗时：{}", request.method(), request.url(), request.headers(), requestBody, System.currentTimeMillis() - startTime))
                .doOnError(e -> log.error("状态：调用大模型异常，请求地址：{}:{}，请求头：{}，参数：{}，耗时：{}", request.method(), request.url(), request.headers(), requestBody, System.currentTimeMillis() - startTime, e));
        return response;
    }

    private ClientRequest buildTraceableRequest(ClientRequest request, StringBuilder requestBody) {
        return ClientRequest.from(request).body((outputMessage, context) -> request.body().insert(new ClientHttpRequestDecorator(outputMessage) {
            @Override
            public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                return super.writeWith(Flux.from(body).doOnNext(buffer -> requestBody.append(buffer.toString(StandardCharsets.UTF_8))));
            }
        }, context)).build();
    }
}
